#define _POSIX_C_SOURCE 200112L
#define _DARWIN_C_SOURCE
#include "storage_engine.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include <pthread.h>
#include <inttypes.h>

// Hash function for mapping table
static uint32_t hash_logical_id(uint64_t logical_id, uint32_t table_size) {
    // Simple hash function for page IDs
    return ((logical_id * 2654435761UL) >> 16) % table_size;
}

// Page data structure functions
static int parse_page_data(bw_page_t* page, char*** keys, uint16_t** key_lengths, 
                           char*** values, uint16_t** value_lengths);
static int pack_page_data(bw_page_t* page, char** keys, uint16_t* key_lengths,
                         char** values, uint16_t* value_lengths, uint16_t count);
static int find_key_in_page_data(bw_page_t* page, const char* key, uint16_t key_len);
static int insert_key_in_page_data(bw_page_t* page, const char* key, uint16_t key_len,
                                  const char* value, uint16_t value_len);

// Get current timestamp
static uint64_t get_current_time(void) {
    struct timespec ts;
    clock_gettime(CLOCK_REALTIME, &ts);
    return (uint64_t)ts.tv_sec * 1000000 + ts.tv_nsec / 1000;
}


// Storage engine initialization
bw_tree_t* storage_init(memory_manager_t* memory_mgr, file_system* fs) {
    bw_tree_t* tree = calloc(1, sizeof(bw_tree_t));
    if (!tree) {
        return NULL;
    }

    // Initialize basic fields
    tree->memory_mgr = memory_mgr;
    tree->fs = fs;
    tree->next_logical_id = 1;
    tree->root_id = 0; // Will be created on first insert
    
    // Configuration
    tree->max_keys_per_page = MAX_KEYS_PER_PAGE;
    tree->min_keys_per_page = MAX_KEYS_PER_PAGE / 2;
    tree->consolidation_threshold = 5;
    
    // Initialize mapping table
    tree->mapping_table_size = 1024;
    tree->mapping_table = calloc(tree->mapping_table_size, sizeof(mapping_entry_t*));
    if (!tree->mapping_table) {
        free(tree);
        return NULL;
    }
    
    // Initialize page cache
    tree->cache_size = 256;
    tree->page_cache = calloc(tree->cache_size, sizeof(bw_page_t*));
    if (!tree->page_cache) {
        free(tree->mapping_table);
        free(tree);
        return NULL;
    }
    
    // Initialize locks
    if (pthread_rwlock_init(&tree->tree_lock, NULL) != 0 ||
        pthread_mutex_init(&tree->id_mutex, NULL) != 0 ||
        pthread_mutex_init(&tree->cache_mutex, NULL) != 0) {
        free(tree->page_cache);
        free(tree->mapping_table);
        free(tree);
        return NULL;
    }
    
    return tree;
}

// Storage engine cleanup
int storage_cleanup(bw_tree_t* tree) {
    if (!tree) {
        return 0;
    }
    
    // Flush any dirty pages
    storage_flush_dirty_pages(tree);
    
    // Cleanup mapping table
    for (uint32_t i = 0; i < tree->mapping_table_size; i++) {
        mapping_entry_t* entry = tree->mapping_table[i];
        while (entry) {
            mapping_entry_t* next = entry->next;
            
            // Cleanup delta chain
            delta_record_t* delta = entry->delta_chain;
            while (delta) {
                delta_record_t* next_delta = delta->next;
                if (delta->type == DELTA_INSERT || delta->type == DELTA_UPDATE || delta->type == DELTA_DELETE) {
                    free(delta->data.modify.kv);
                }
                free(delta);
                delta = next_delta;
            }
            
            pthread_rwlock_destroy(&entry->lock);
            free(entry);
            entry = next;
        }
    }
    
    // Cleanup page cache
    pthread_mutex_lock(&tree->cache_mutex);
    for (uint32_t i = 0; i < tree->cache_count; i++) {
        free(tree->page_cache[i]);
    }
    pthread_mutex_unlock(&tree->cache_mutex);
    
    // Destroy locks
    pthread_rwlock_destroy(&tree->tree_lock);
    pthread_mutex_destroy(&tree->id_mutex);
    pthread_mutex_destroy(&tree->cache_mutex);
    
    // Free memory
    free(tree->mapping_table);
    free(tree->page_cache);
    free(tree);
    
    return 0;
}

// Generate unique logical ID
uint64_t storage_generate_logical_id(bw_tree_t* tree) {
    pthread_mutex_lock(&tree->id_mutex);
    uint64_t id = tree->next_logical_id++;
    pthread_mutex_unlock(&tree->id_mutex);
    return id;
}

// Key comparison function
int storage_compare_keys(const char* key1, uint16_t len1, const char* key2, uint16_t len2) {
    uint16_t min_len = len1 < len2 ? len1 : len2;
    int result = memcmp(key1, key2, min_len);
    
    if (result == 0) {
        if (len1 < len2) return KEY_LESS;
        if (len1 > len2) return KEY_GREATER;
        return KEY_EQUAL;
    }
    
    return result < 0 ? KEY_LESS : KEY_GREATER;
}

// Create key-value pair
kv_pair_t* storage_create_kv_pair(const char* key, uint16_t key_len, 
                                 const char* value, uint16_t value_len) {
    size_t total_size = sizeof(kv_pair_t) + key_len + value_len;
    kv_pair_t* kv = malloc(total_size);
    if (!kv) {
        return NULL;
    }
    
    kv->key_len = key_len;
    kv->value_len = value_len;
    memcpy(kv->data, key, key_len);
    memcpy(kv->data + key_len, value, value_len);
    
    return kv;
}

// Get mapping entry for logical ID
mapping_entry_t* storage_get_mapping(bw_tree_t* tree, uint64_t logical_id) {
    uint32_t hash = hash_logical_id(logical_id, tree->mapping_table_size);
    mapping_entry_t* entry = tree->mapping_table[hash];
    
    while (entry) {
        if (entry->logical_id == logical_id) {
            return entry;
        }
        entry = entry->next;
    }
    
    return NULL;
}

// Update mapping table
int storage_update_mapping(bw_tree_t* tree, uint64_t logical_id, uint64_t physical_id) {
    uint32_t hash = hash_logical_id(logical_id, tree->mapping_table_size);
    mapping_entry_t* entry = storage_get_mapping(tree, logical_id);
    
    if (entry) {
        // Update existing mapping
        pthread_rwlock_wrlock(&entry->lock);
        entry->physical_id = physical_id;
        pthread_rwlock_unlock(&entry->lock);
        return 0;
    }
    
    // Create new mapping entry
    entry = calloc(1, sizeof(mapping_entry_t));
    if (!entry) {
        return -1;
    }
    
    entry->logical_id = logical_id;
    entry->physical_id = physical_id;
    entry->delta_chain = NULL;
    
    if (pthread_rwlock_init(&entry->lock, NULL) != 0) {
        free(entry);
        return -1;
    }
    
    // Add to hash table
    entry->next = tree->mapping_table[hash];
    tree->mapping_table[hash] = entry;
    tree->mapping_count++;
    
    return 0;
}

// Add delta record to mapping
int storage_add_delta(bw_tree_t* tree, uint64_t logical_id, delta_record_t* delta) {
    mapping_entry_t* entry = storage_get_mapping(tree, logical_id);
    if (!entry) {
        return -1;
    }
    
    pthread_rwlock_wrlock(&entry->lock);
    
    // Add delta to front of chain
    delta->next = entry->delta_chain;
    entry->delta_chain = delta;
    tree->delta_count++;
    
    pthread_rwlock_unlock(&entry->lock);
    
    // Check if consolidation is needed
    int delta_count = 0;
    delta_record_t* d = entry->delta_chain;
    while (d && delta_count < tree->consolidation_threshold) {
        delta_count++;
        d = d->next;
    }
    
    if (delta_count >= tree->consolidation_threshold) {
        // Schedule consolidation (for now, do it immediately)
        storage_consolidate_page(tree, logical_id);
    }
    
    return 0;
}

// Load page from storage
bw_page_t* storage_load_page(bw_tree_t* tree, uint64_t physical_id) {
    // Check cache first
    pthread_mutex_lock(&tree->cache_mutex);
    for (uint32_t i = 0; i < tree->cache_count; i++) {
        if (tree->page_cache[i] && tree->page_cache[i]->page_id == physical_id) {
            bw_page_t* page = tree->page_cache[i];
            pthread_mutex_unlock(&tree->cache_mutex);
            return page;
        }
    }
    pthread_mutex_unlock(&tree->cache_mutex);
    
    // Load from file system
    bw_page_t* page = malloc(sizeof(bw_page_t));
    if (!page) {
        return NULL;
    }
    
    // For now, create a dummy page (in real implementation, read from file)
    memset(page, 0, sizeof(bw_page_t));
    page->page_id = physical_id;
    page->node_type = NODE_TYPE_LEAF;
    page->key_count = 0;
    page->free_space = PAGE_SIZE - 32;
    page->is_dirty = false;
    page->lsn = 0;
    page->timestamp = get_current_time();
    
    // Add to cache if there's space
    pthread_mutex_lock(&tree->cache_mutex);
    if (tree->cache_count < tree->cache_size) {
        tree->page_cache[tree->cache_count++] = page;
    }
    pthread_mutex_unlock(&tree->cache_mutex);
    
    return page;
}

// Store page to storage
int storage_store_page(bw_tree_t* tree, bw_page_t* page) {
    if (!tree || !page) {
        return -1;
    }
    
    // Mark page as clean
    page->is_dirty = false;
    page->timestamp = get_current_time();
    
    // Update statistics
    tree->total_pages++;
    if (page->node_type == NODE_TYPE_LEAF) {
        tree->leaf_pages++;
    } else {
        tree->inner_pages++;
    }
    
    return 0;
}

// Consolidate page with delta chain
bw_node_t* storage_consolidate_page(bw_tree_t* tree, uint64_t logical_id) {
    mapping_entry_t* entry = storage_get_mapping(tree, logical_id);
    if (!entry) {
        return NULL;
    }
    
    pthread_rwlock_wrlock(&entry->lock);
    
    // Load base page
    bw_page_t* base_page = storage_load_page(tree, entry->physical_id);
    if (!base_page) {
        pthread_rwlock_unlock(&entry->lock);
        return NULL;
    }
    
    // Create consolidated node
    bw_node_t* node = calloc(1, sizeof(bw_node_t));
    if (!node) {
        pthread_rwlock_unlock(&entry->lock);
        return NULL;
    }
    
    node->logical_id = logical_id;
    node->node_type = base_page->node_type;
    node->key_count = base_page->key_count;
    node->is_consolidated = true;
    node->epoch = get_current_time();
    
    // Allocate arrays
    node->keys = calloc(tree->max_keys_per_page, sizeof(char*));
    node->key_lengths = calloc(tree->max_keys_per_page, sizeof(uint16_t));
    
    if (node->node_type == NODE_TYPE_LEAF) {
        node->data.leaf.values = calloc(tree->max_keys_per_page, sizeof(char*));
        node->data.leaf.value_lengths = calloc(tree->max_keys_per_page, sizeof(uint16_t));
    } else {
        node->data.inner.children = calloc(tree->max_keys_per_page + 1, sizeof(uint64_t));
    }
    
    // Apply delta chain to create consolidated view
    char** keys = NULL;
    uint16_t* key_lengths = NULL;
    char** values = NULL;
    uint16_t* value_lengths = NULL;
    
    // Parse base page data
    int base_count = parse_page_data(base_page, &keys, &key_lengths, &values, &value_lengths);
    if (base_count < 0) {
        base_count = 0;
    }
    
    // Apply deltas in reverse order (from oldest to newest)
    delta_record_t* deltas[100]; // Maximum delta chain length
    int delta_count = 0;
    delta_record_t* delta = entry->delta_chain;
    while (delta && delta_count < 100) {
        deltas[delta_count++] = delta;
        delta = delta->next;
    }
    
    // Apply deltas from oldest to newest
    for (int i = delta_count - 1; i >= 0; i--) {
        delta = deltas[i];
        kv_pair_t* kv = delta->data.modify.kv;
        
        if (delta->type == DELTA_INSERT || delta->type == DELTA_UPDATE) {
            // Find if key already exists
            int existing_pos = -1;
            for (int j = 0; j < base_count; j++) {
                if (storage_compare_keys(keys[j], key_lengths[j], kv->data, kv->key_len) == KEY_EQUAL) {
                    existing_pos = j;
                    break;
                }
            }
            
            if (existing_pos >= 0) {
                // Update existing key
                free(values[existing_pos]);
                values[existing_pos] = malloc(kv->value_len);
                if (values[existing_pos]) {
                    memcpy(values[existing_pos], kv->data + kv->key_len, kv->value_len);
                    value_lengths[existing_pos] = kv->value_len;
                }
            } else if (base_count < tree->max_keys_per_page) {
                // Insert new key
                keys[base_count] = malloc(kv->key_len);
                values[base_count] = malloc(kv->value_len);
                if (keys[base_count] && values[base_count]) {
                    memcpy(keys[base_count], kv->data, kv->key_len);
                    memcpy(values[base_count], kv->data + kv->key_len, kv->value_len);
                    key_lengths[base_count] = kv->key_len;
                    value_lengths[base_count] = kv->value_len;
                    base_count++;
                }
            }
        } else if (delta->type == DELTA_DELETE) {
            // Find and remove key
            for (int j = 0; j < base_count; j++) {
                if (storage_compare_keys(keys[j], key_lengths[j], kv->data, kv->key_len) == KEY_EQUAL) {
                    // Remove key-value pair
                    free(keys[j]);
                    free(values[j]);
                    // Shift remaining entries
                    for (int k = j; k < base_count - 1; k++) {
                        keys[k] = keys[k + 1];
                        values[k] = values[k + 1];
                        key_lengths[k] = key_lengths[k + 1];
                        value_lengths[k] = value_lengths[k + 1];
                    }
                    base_count--;
                    break;
                }
            }
        }
    }
    
    // Update node with consolidated data
    node->key_count = base_count;
    if (keys && base_count > 0) {
        for (int i = 0; i < base_count; i++) {
            node->keys[i] = keys[i];
            node->key_lengths[i] = key_lengths[i];
            if (node->node_type == NODE_TYPE_LEAF) {
                node->data.leaf.values[i] = values[i];
                node->data.leaf.value_lengths[i] = value_lengths[i];
            }
        }
    }
    
    // Cleanup arrays (but not the data, which is now owned by the node)
    if (keys) free(keys);
    if (key_lengths) free(key_lengths);
    if (values) free(values);
    if (value_lengths) free(value_lengths);
    
    // Clear delta chain
    delta_record_t* old_delta = entry->delta_chain;
    while (old_delta) {
        delta_record_t* next = old_delta->next;
        if (old_delta->type == DELTA_INSERT || old_delta->type == DELTA_UPDATE || old_delta->type == DELTA_DELETE) {
            free(old_delta->data.modify.kv);
        }
        free(old_delta);
        old_delta = next;
    }
    entry->delta_chain = NULL;
    tree->consolidation_count++;
    
    pthread_rwlock_unlock(&entry->lock);
    
    return node;
}

// Search for a key in the tree
search_result_t storage_search(bw_tree_t* tree, const char* key, uint16_t key_len) {
    search_result_t result = {false, NULL, 0, 0, 0};
    
    if (!tree || !key) {
        return result;
    }
    
    // Avoid unused parameter warning
    (void)key_len;
    
    pthread_rwlock_rdlock(&tree->tree_lock);
    
    // Handle empty tree
    if (tree->root_id == 0) {
        pthread_rwlock_unlock(&tree->tree_lock);
        return result;
    }
    
    uint64_t current_id = tree->root_id;
    
    // Traverse tree to find leaf page
    while (true) {
        mapping_entry_t* entry = storage_get_mapping(tree, current_id);
        if (!entry) {
            break;
        }
        
        pthread_rwlock_rdlock(&entry->lock);
        bw_page_t* page = storage_load_page(tree, entry->physical_id);
        
        if (!page) {
            pthread_rwlock_unlock(&entry->lock);
            break;
        }
        
        if (page->node_type == NODE_TYPE_LEAF) {
            // Search in leaf page
            result.logical_id = current_id;
            
            // Apply deltas to get current view
            char** keys = NULL;
            uint16_t* key_lengths = NULL;
            char** values = NULL;
            uint16_t* value_lengths = NULL;
            
            if (parse_page_data(page, &keys, &key_lengths, &values, &value_lengths) >= 0) {
                // Apply delta chain
                delta_record_t* delta = entry->delta_chain;
                while (delta) {
                    if (delta->type == DELTA_INSERT || delta->type == DELTA_UPDATE) {
                        kv_pair_t* kv = delta->data.modify.kv;
                        // Check if this key matches
                        if (storage_compare_keys(kv->data, kv->key_len, key, key_len) == KEY_EQUAL) {
                            result.found = true;
                            if (result.value) free(result.value);
                            result.value = malloc(kv->value_len);
                            if (result.value) {
                                memcpy(result.value, kv->data + kv->key_len, kv->value_len);
                                result.value_len = kv->value_len;
                            }
                            break;
                        }
                    } else if (delta->type == DELTA_DELETE) {
                        kv_pair_t* kv = delta->data.modify.kv;
                        if (storage_compare_keys(kv->data, kv->key_len, key, key_len) == KEY_EQUAL) {
                            result.found = false;
                            break;
                        }
                    }
                    delta = delta->next;
                }
                
                // Check base page if not found in deltas
                if (!result.found) {
                    int pos = find_key_in_page_data(page, key, key_len);
                    if (pos >= 0) {
                        result.found = true;
                        result.position = pos;
                        if (values && values[pos]) {
                            result.value = malloc(value_lengths[pos]);
                            if (result.value) {
                                memcpy(result.value, values[pos], value_lengths[pos]);
                                result.value_len = value_lengths[pos];
                            }
                        }
                    }
                }
                
                // Cleanup
                if (keys) {
                    for (int i = 0; i < page->key_count; i++) {
                        free(keys[i]);
                        if (values) free(values[i]);
                    }
                    free(keys);
                    free(key_lengths);
                    if (values) free(values);
                    if (value_lengths) free(value_lengths);
                }
            }
            
            pthread_rwlock_unlock(&entry->lock);
            break;
        } else {
            // Inner node - find child to descend to
            // For now, just go to first child (simplified implementation)
            current_id = 1; // Placeholder - should be actual child ID
            
            pthread_rwlock_unlock(&entry->lock);
        }
    }
    
    pthread_rwlock_unlock(&tree->tree_lock);
    return result;
}

// Insert key-value pair
int storage_insert(bw_tree_t* tree, const char* key, uint16_t key_len, 
                  const char* value, uint16_t value_len) {
    if (!tree || !key || !value) {
        return -1;
    }
    
    pthread_rwlock_wrlock(&tree->tree_lock);
    
    // Create root if tree is empty
    if (tree->root_id == 0) {
        tree->root_id = storage_generate_logical_id(tree);
        
        // Create root page
        bw_page_t* root_page = malloc(sizeof(bw_page_t));
        if (!root_page) {
            pthread_rwlock_unlock(&tree->tree_lock);
            return -1;
        }
        
        memset(root_page, 0, sizeof(bw_page_t));
        root_page->page_id = tree->root_id;
        root_page->node_type = NODE_TYPE_LEAF;
        root_page->key_count = 0;
        root_page->free_space = PAGE_SIZE - 32;
        root_page->is_dirty = true;
        root_page->timestamp = get_current_time();
        
        storage_store_page(tree, root_page);
        storage_update_mapping(tree, tree->root_id, tree->root_id);
    }
    
    // Find target leaf page
    search_result_t search = storage_search(tree, key, key_len);
    
    // Create insert delta
    delta_record_t* delta = malloc(sizeof(delta_record_t));
    if (!delta) {
        pthread_rwlock_unlock(&tree->tree_lock);
        return -1;
    }
    
    delta->type = DELTA_INSERT;
    delta->logical_id = search.logical_id ? search.logical_id : tree->root_id;
    delta->data.modify.kv = storage_create_kv_pair(key, key_len, value, value_len);
    
    if (!delta->data.modify.kv) {
        free(delta);
        pthread_rwlock_unlock(&tree->tree_lock);
        return -1;
    }
    
    // Add delta to mapping
    int result = storage_add_delta(tree, delta->logical_id, delta);
    
    pthread_rwlock_unlock(&tree->tree_lock);
    return result;
}

// Delete key from tree
int storage_delete(bw_tree_t* tree, const char* key, uint16_t key_len) {
    if (!tree || !key) {
        return -1;
    }
    
    pthread_rwlock_wrlock(&tree->tree_lock);
    
    // Find key location
    search_result_t search = storage_search(tree, key, key_len);
    if (!search.found) {
        pthread_rwlock_unlock(&tree->tree_lock);
        return -1; // Key not found
    }
    
    // Create delete delta
    delta_record_t* delta = malloc(sizeof(delta_record_t));
    if (!delta) {
        pthread_rwlock_unlock(&tree->tree_lock);
        return -1;
    }
    
    delta->type = DELTA_DELETE;
    delta->logical_id = search.logical_id;
    delta->data.modify.kv = storage_create_kv_pair(key, key_len, "", 0);
    
    if (!delta->data.modify.kv) {
        free(delta);
        pthread_rwlock_unlock(&tree->tree_lock);
        return -1;
    }
    
    // Add delta to mapping
    int result = storage_add_delta(tree, delta->logical_id, delta);
    
    pthread_rwlock_unlock(&tree->tree_lock);
    return result;
}

// Update existing key with new value
int storage_update(bw_tree_t* tree, const char* key, uint16_t key_len, 
                  const char* value, uint16_t value_len) {
    if (!tree || !key || !value) {
        return -1;
    }
    
    // For Bw-Tree, update is implemented as delete + insert
    // First check if key exists
    search_result_t search = storage_search(tree, key, key_len);
    if (!search.found) {
        return -1; // Key not found
    }
    
    pthread_rwlock_wrlock(&tree->tree_lock);
    
    // Create update delta
    delta_record_t* delta = malloc(sizeof(delta_record_t));
    if (!delta) {
        pthread_rwlock_unlock(&tree->tree_lock);
        return -1;
    }
    
    delta->type = DELTA_UPDATE;
    delta->logical_id = search.logical_id;
    delta->data.modify.kv = storage_create_kv_pair(key, key_len, value, value_len);
    
    if (!delta->data.modify.kv) {
        free(delta);
        pthread_rwlock_unlock(&tree->tree_lock);
        return -1;
    }
    
    // Add delta to mapping
    int result = storage_add_delta(tree, delta->logical_id, delta);
    
    pthread_rwlock_unlock(&tree->tree_lock);
    return result;
}

// Flush dirty pages to storage
int storage_flush_dirty_pages(bw_tree_t* tree) {
    if (!tree) {
        return -1;
    }
    
    pthread_mutex_lock(&tree->cache_mutex);
    
    int flushed = 0;
    for (uint32_t i = 0; i < tree->cache_count; i++) {
        if (tree->page_cache[i] && tree->page_cache[i]->is_dirty) {
            storage_store_page(tree, tree->page_cache[i]);
            flushed++;
        }
    }
    
    pthread_mutex_unlock(&tree->cache_mutex);
    return flushed;
}

// Print storage engine statistics
void storage_print_stats(const bw_tree_t* tree) {
    if (!tree) {
        return;
    }
    
    printf("Bw-Tree Storage Engine Statistics:\n");
    printf("  Root ID: %" PRIu64 "\n", tree->root_id);
    printf("  Next Logical ID: %" PRIu64 "\n", tree->next_logical_id);
    printf("  Total Pages: %" PRIu64 "\n", tree->total_pages);
    printf("  Leaf Pages: %" PRIu64 "\n", tree->leaf_pages);
    printf("  Inner Pages: %" PRIu64 "\n", tree->inner_pages);
    printf("  Mapping Count: %u\n", tree->mapping_count);
    printf("  Delta Records: %" PRIu64 "\n", tree->delta_count);
    printf("  Split Operations: %" PRIu64 "\n", tree->split_count);
    printf("  Merge Operations: %" PRIu64 "\n", tree->merge_count);
    printf("  Consolidations: %" PRIu64 "\n", tree->consolidation_count);
    printf("  Cache Size: %u/%u\n", tree->cache_count, tree->cache_size);
    printf("  Max Keys per Page: %u\n", tree->max_keys_per_page);
    printf("  Min Keys per Page: %u\n", tree->min_keys_per_page);
    printf("  Consolidation Threshold: %u\n", tree->consolidation_threshold);
}

// Parse page data into separate arrays
static int parse_page_data(bw_page_t* page, char*** keys, uint16_t** key_lengths, 
                           char*** values, uint16_t** value_lengths) {
    if (!page || !keys || !key_lengths || !values || !value_lengths) {
        return -1;
    }
    
    if (page->key_count == 0) {
        *keys = NULL;
        *key_lengths = NULL;
        *values = NULL;
        *value_lengths = NULL;
        return 0;
    }
    
    // Allocate arrays
    *keys = calloc(page->key_count, sizeof(char*));
    *key_lengths = calloc(page->key_count, sizeof(uint16_t));
    *values = calloc(page->key_count, sizeof(char*));
    *value_lengths = calloc(page->key_count, sizeof(uint16_t));
    
    if (!*keys || !*key_lengths || !*values || !*value_lengths) {
        free(*keys);
        free(*key_lengths);
        free(*values);
        free(*value_lengths);
        return -1;
    }
    
    // Parse page data - simple format: [key_len][key][value_len][value]...
    char* data_ptr = page->data;
    char* data_end = page->data + sizeof(page->data);
    
    for (int i = 0; i < page->key_count && data_ptr < data_end; i++) {
        // Read key length
        if (data_ptr + sizeof(uint16_t) > data_end) break;
        (*key_lengths)[i] = *(uint16_t*)data_ptr;
        data_ptr += sizeof(uint16_t);
        
        // Read key
        if (data_ptr + (*key_lengths)[i] > data_end) break;
        (*keys)[i] = malloc((*key_lengths)[i]);
        if ((*keys)[i]) {
            memcpy((*keys)[i], data_ptr, (*key_lengths)[i]);
        }
        data_ptr += (*key_lengths)[i];
        
        // Read value length
        if (data_ptr + sizeof(uint16_t) > data_end) break;
        (*value_lengths)[i] = *(uint16_t*)data_ptr;
        data_ptr += sizeof(uint16_t);
        
        // Read value
        if (data_ptr + (*value_lengths)[i] > data_end) break;
        (*values)[i] = malloc((*value_lengths)[i]);
        if ((*values)[i]) {
            memcpy((*values)[i], data_ptr, (*value_lengths)[i]);
        }
        data_ptr += (*value_lengths)[i];
    }
    
    return page->key_count;
}

// Pack key-value arrays into page data
static int pack_page_data(bw_page_t* page, char** keys, uint16_t* key_lengths,
                         char** values, uint16_t* value_lengths, uint16_t count) {
    if (!page) return -1;
    
    if (count == 0) {
        page->key_count = 0;
        page->free_space = sizeof(page->data);
        return 0;
    }
    
    char* data_ptr = page->data;
    char* data_end = page->data + sizeof(page->data);
    uint16_t packed_count = 0;
    
    for (int i = 0; i < count && data_ptr < data_end; i++) {
        size_t needed_space = sizeof(uint16_t) + key_lengths[i] + sizeof(uint16_t) + value_lengths[i];
        if (data_ptr + needed_space > data_end) {
            break; // Page full
        }
        
        // Write key length
        *(uint16_t*)data_ptr = key_lengths[i];
        data_ptr += sizeof(uint16_t);
        
        // Write key
        memcpy(data_ptr, keys[i], key_lengths[i]);
        data_ptr += key_lengths[i];
        
        // Write value length
        *(uint16_t*)data_ptr = value_lengths[i];
        data_ptr += sizeof(uint16_t);
        
        // Write value
        memcpy(data_ptr, values[i], value_lengths[i]);
        data_ptr += value_lengths[i];
        
        packed_count++;
    }
    
    page->key_count = packed_count;
    page->free_space = data_end - data_ptr;
    page->is_dirty = true;
    
    return packed_count;
}

// Find key in page data, returns position or -1 if not found
static int find_key_in_page_data(bw_page_t* page, const char* key, uint16_t key_len) {
    if (!page || !key || page->key_count == 0) {
        return -1;
    }
    
    char* data_ptr = page->data;
    char* data_end = page->data + sizeof(page->data);
    
    for (int i = 0; i < page->key_count && data_ptr < data_end; i++) {
        // Read key length
        if (data_ptr + sizeof(uint16_t) > data_end) break;
        uint16_t stored_key_len = *(uint16_t*)data_ptr;
        data_ptr += sizeof(uint16_t);
        
        // Compare key
        if (data_ptr + stored_key_len > data_end) break;
        if (storage_compare_keys(data_ptr, stored_key_len, key, key_len) == KEY_EQUAL) {
            return i;
        }
        data_ptr += stored_key_len;
        
        // Skip value length and value
        if (data_ptr + sizeof(uint16_t) > data_end) break;
        uint16_t value_len = *(uint16_t*)data_ptr;
        data_ptr += sizeof(uint16_t) + value_len;
    }
    
    return -1;
}

// Insert key-value pair into page data
static int insert_key_in_page_data(bw_page_t* page, const char* key, uint16_t key_len,
                                  const char* value, uint16_t value_len) {
    if (!page || !key || !value) {
        return -1;
    }
    
    size_t needed_space = sizeof(uint16_t) + key_len + sizeof(uint16_t) + value_len;
    if (page->free_space < needed_space) {
        return -1; // Page full
    }
    
    // Find insertion point (maintain sorted order)
    char* insert_point = page->data;
    char* data_ptr = page->data;
    char* data_end = page->data + sizeof(page->data);
    int insert_pos = 0;
    
    for (int i = 0; i < page->key_count && data_ptr < data_end; i++) {
        // Read key length
        if (data_ptr + sizeof(uint16_t) > data_end) break;
        uint16_t stored_key_len = *(uint16_t*)data_ptr;
        data_ptr += sizeof(uint16_t);
        
        // Compare key
        if (data_ptr + stored_key_len > data_end) break;
        int cmp = storage_compare_keys(data_ptr, stored_key_len, key, key_len);
        if (cmp >= 0) {
            insert_point = data_ptr - sizeof(uint16_t);
            break;
        }
        data_ptr += stored_key_len;
        
        // Skip value
        if (data_ptr + sizeof(uint16_t) > data_end) break;
        uint16_t value_len_stored = *(uint16_t*)data_ptr;
        data_ptr += sizeof(uint16_t) + value_len_stored;
        insert_pos++;
        insert_point = data_ptr;
    }
    
    // If we need to insert in the middle, shift data
    if (insert_pos < page->key_count) {
        size_t remaining_data = (page->data + sizeof(page->data) - page->free_space) - insert_point;
        memmove(insert_point + needed_space, insert_point, remaining_data);
    }
    
    // Insert the new key-value pair
    char* write_ptr = insert_point;
    *(uint16_t*)write_ptr = key_len;
    write_ptr += sizeof(uint16_t);
    memcpy(write_ptr, key, key_len);
    write_ptr += key_len;
    *(uint16_t*)write_ptr = value_len;
    write_ptr += sizeof(uint16_t);
    memcpy(write_ptr, value, value_len);
    
    page->key_count++;
    page->free_space -= needed_space;
    page->is_dirty = true;
    
    return insert_pos;
}